import paddle
import cv2
import numpy as np
from ppdet.modeling.architectures import YOLOv3
from ppdet.modeling.backbones import mobilenet_v1
from ppdet.core.workspace import load_config, create
from ppdet.utils.download import download_dataset

def init_detector():
    # 创建YOLOv3模型实例
    model = YOLOv3(
        backbone=mobilenet_v1(),
        num_classes=80  # COCO数据集的类别数
    )
    
    # 加载预训练权重
    #weights_url = 'https://paddledet.bj.bcebos.com/models/yolov3_mobilenet_v1_270e_coco.pdparams'
    #weights_path = download_dataset(weights_url)
    weights_path = 'yolov3_mobilenet_v1_270e_coco.pdparams'
    model.set_state_dict(paddle.load(weights_path))
    model.eval()
    return model

def detect_image(image_path):
    # 读取图像
    img = cv2.imread(image_path)
    if img is None:
        raise ValueError('无法读取图像，请检查图像路径')
    
    # 图像预处理
    resized_img = cv2.resize(img, (608, 608))
    img_rgb = cv2.cvtColor(resized_img, cv2.COLOR_BGR2RGB)
    img_normalized = img_rgb.astype('float32') / 255.
    img_transposed = np.transpose(img_normalized, (2, 0, 1))
    img_expanded = np.expand_dims(img_transposed, axis=0)
    
    # 创建检测器并加载权重
    model = init_detector()
    
    # 执行检测
    with paddle.no_grad():
        input_tensor = paddle.to_tensor(img_expanded)
        results = model(input_tensor)
    
    # 处理检测结果
    threshold = 0.5
    for box, score, label in zip(results[0][0], results[1][0], results[2][0]):
        if score < threshold:
            continue
        
        x1, y1, x2, y2 = box.numpy().astype('int')
        cv2.rectangle(img, (x1, y1), (x2, y2), (0, 255, 0), 2)
        text = f'Class: {int(label.numpy())}, Score: {score.numpy():.2f}'
        cv2.putText(img, text, (x1, y1-10), cv2.FONT_HERSHEY_SIMPLEX, 0.5, (0, 255, 0), 2)
    
    # 保存结果
    cv2.imwrite('detection_result.jpg', img)
    print('检测结果已保存为 detection_result.jpg')

if __name__ == '__main__':
    # 替换为你的测试图片路径
    image_path = 'test.png'
    detect_image(image_path)